import os

import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope
# 0.12.48 版本的正确导入方式
from llama_index.vector_stores.chroma import ChromaVectorStore  # 修改这一行

data_dir = 'D:/pythonRoom/llm_llama/data'
# 配置Chroma向量数据库
PERSIST_DIR = "D:/pythonRoom/llm_llama/chroma_db"  # 向量数据库存储目录
COLLECTION_NAME = "knowledge_base"  # 集合名称
'''
阿里百炼地址
https://bailian.console.aliyun.com/?tab=app#/api-key
'''
# 1. 配置阿里云API（替换成您的API Key）
print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

# 2. 使用本地嵌入模型（节省成本）
print('使用本地嵌入模型（节省成本）')
embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)

# 3. 初始化Chroma客户端
chroma_client = chromadb.PersistentClient(path=PERSIST_DIR)
chroma_collection = chroma_client.get_or_create_collection(COLLECTION_NAME)

# 4. 创建向量存储
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 5. 检查向量数据库是否已有数据
if chroma_collection.count() == 0:
    # 5.1 加载文档
    print('加载文档')
    documents = SimpleDirectoryReader(data_dir).load_data()

    # 5.2 构建索引
    print('构建索引')
    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        embed_model=embed_model
    )

    print(f'向量数据库已创建，包含 {chroma_collection.count()} 条向量')
else:
    # 5.3 从向量数据库加载已有索引
    print(f'从向量数据库加载已有索引，包含 {chroma_collection.count()} 条向量')
    index = VectorStoreIndex.from_vector_store(
        vector_store=vector_store,
        embed_model=embed_model  # 必须使用相同的嵌入模型
    )

# 6. 创建查询引擎
print('创建查询引擎')
query_engine = index.as_query_engine(llm=llm)

# 7. 提问
print('提问')
response = query_engine.query("李大长是如何发家致富的?")
print("回答:", response)


# 8. 添加文档更新功能示例
def add_document_to_knowledge_base(file_path: str):
    """向知识库添加新文档"""
    print(f"添加新文档: {file_path}")

    # 加载新文档
    new_documents = SimpleDirectoryReader(input_files=[file_path]).load_data()

    # 创建新索引器
    new_index = VectorStoreIndex.from_documents(
        new_documents,
        storage_context=storage_context,
        embed_model=embed_model
    )

    print(f"文档已添加，当前向量总数: {chroma_collection.count()}")

# 示例：添加新文档
# add_document_to_knowledge_base("new_data/document.pdf")
